package com.example.yckjbigdataflink.job;

import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.datastream.AsyncDataStream;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

import org.apache.flink.streaming.api.functions.async.ResultFuture;
import org.apache.flink.streaming.api.functions.async.RichAsyncFunction;
import org.apache.flink.streaming.api.functions.source.SourceFunction;
import org.apache.flink.table.api.*;
import org.apache.flink.table.api.bridge.java.StreamTableEnvironment;

import java.util.*;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.TimeUnit;

import static org.apache.flink.table.api.Expressions.$;

/**
 * Flink批处理示范：
 * 1. 切分RowKey区间
 * 2. 异步调用HBase读取模拟（请替换为真实HBase客户端调用）
 * 3. 将结果写入Hive表
 */
public class FlinkHBaseToHiveJob {

    public static void main(String[] args) throws Exception {
        // 1. 环境配置
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(5);

        StreamTableEnvironment tableEnv = StreamTableEnvironment.create(env);

        // 2. 生成RowKey切片
        List<RowKeyRange> rowKeyRanges = RowKeySplitter.splitByPrefix(5);

        // 3. 通过Source发射RowKey区间（每个范围触发一次读取）
        DataStream<RowKeyRange> sourceStream = env.addSource(new RowKeyRangeSource(rowKeyRanges));

        // 4. 异步调用HBase读取数据，模拟异步查询
        DataStream<HBaseResult> hbaseResultStream = AsyncDataStream.unorderedWait(
                sourceStream,
                new AsyncHBaseReader(),
                10_000, // 超时10秒
                TimeUnit.MILLISECONDS,
                20 // 最大并发20
        );

        // 5. 转成Table写入Hive
        Table resultTable = tableEnv.fromDataStream(hbaseResultStream,
                $("rowkey"), $("col1"), $("col2"), $("dt"));

        // 6. 注册Hive Catalog并切换到对应库（请根据实际调整）
        tableEnv.executeSql(
                "CREATE CATALOG myhive WITH (" +
                        "'type'='hive'," +
                        "'hive-conf-dir'='/path/to/hive/conf')"  // 修改为真实Hive配置路径
        );
        tableEnv.useCatalog("myhive");
        tableEnv.useDatabase("mydb");

        // 7. 创建Hive表（如果未创建）
        tableEnv.executeSql(
                "CREATE TABLE IF NOT EXISTS processed_data (" +
                        "rowkey STRING," +
                        "col1 STRING," +
                        "col2 INT," +
                        "dt STRING" +
                        ") PARTITIONED BY (dt) STORED AS PARQUET"
        );

        // 8. 写入Hive（动态分区）
        resultTable.executeInsert("processed_data");

        env.execute("Flink HBase Async Read to Hive Sink Job");
    }

    // -------------------- 工具类 --------------------

    /** RowKey区间表示 */
    public static class RowKeyRange {
        public String startRowKey;
        public String endRowKey;

        public RowKeyRange() {
        }

        public RowKeyRange(String startRowKey, String endRowKey) {
            this.startRowKey = startRowKey;
            this.endRowKey = endRowKey;
        }
    }

    /** RowKey切片生成工具 */
    public static class RowKeySplitter {
        public static List<RowKeyRange> splitByPrefix(int numShards) {
            List<RowKeyRange> shards = new ArrayList<>();
            char start = '0';
            char end = 'z';

            int totalRange = end - start + 1;
            int shardSize = totalRange / numShards;

            for (int i = 0; i < numShards; i++) {
                char shardStart = (char) (start + i * shardSize);
                char shardEnd = (i == numShards - 1) ? end : (char) (shardStart + shardSize - 1);

                String startKey = String.valueOf(shardStart);
                String endKey = String.valueOf(shardEnd) + "~";

                shards.add(new RowKeyRange(startKey, endKey));
            }
            return shards;
        }
    }

    /** 用于发送RowKeyRange的Source */
    public static class RowKeyRangeSource implements SourceFunction<RowKeyRange> {
        private final List<RowKeyRange> ranges;
        private volatile boolean running = true;

        public RowKeyRangeSource(List<RowKeyRange> ranges) {
            this.ranges = ranges;
        }

        @Override
        public void run(SourceContext<RowKeyRange> ctx) throws Exception {
            for (RowKeyRange range : ranges) {
                if (!running) break;
                ctx.collect(range);
                Thread.sleep(100); // 稍微间隔下，防止压力
            }
        }

        @Override
        public void cancel() {
            running = false;
        }
    }

    /** 模拟HBase异步读取的结果实体 */
    public static class HBaseResult {
        public String rowkey;
        public String col1;
        public Integer col2;
        public String dt;

        public HBaseResult() {}

        public HBaseResult(String rowkey, String col1, Integer col2, String dt) {
            this.rowkey = rowkey;
            this.col1 = col1;
            this.col2 = col2;
            this.dt = dt;
        }
    }

    /** 异步读取HBase的示范函数，需替换成真实HBase异步API调用 */
    public static class AsyncHBaseReader extends RichAsyncFunction<RowKeyRange, HBaseResult> {

        @Override
        public void open(Configuration parameters) throws Exception {
            super.open(parameters);
            // 初始化HBase客户端、Kerberos认证等
        }

        @Override
        public void asyncInvoke(RowKeyRange input, ResultFuture<HBaseResult> resultFuture) {
            // TODO: 这里用真实异步HBase客户端调用扫描数据，下面是模拟数据示范

            // 模拟异步执行：返回10条数据
            CompletableFuture.runAsync(() -> {
                List<HBaseResult> resultList = new ArrayList<>();
                for (int i = 0; i < 10; i++) {
                    String rk = input.startRowKey + "_row_" + i;
                    String col1 = "value_" + i;
                    int col2 = i * 10;
                    String dt = java.time.LocalDate.now().toString();
                    resultList.add(new HBaseResult(rk, col1, col2, dt));
                }
                resultFuture.complete(resultList);
            });
        }

        @Override
        public void timeout(RowKeyRange input, ResultFuture<HBaseResult> resultFuture) {
            // 超时处理逻辑：可以重试或返回空列表
            resultFuture.complete(Collections.emptyList());
        }
    }
}
